import torch
from model.SAC.v2.sac_learning import SACContinuous
from environment.electric_scheduling import PowerDayAheadSchedule

class DayAheadScheduleModel:
    def __init__(self, env: PowerDayAheadSchedule, cfg):
        super(DayAheadScheduleModel, self).__init__()
        self.device = cfg['device']
        self.fire_model = SACContinuous(
            env,
            state_dim=cfg["state_dim"],
            hidden_dim=cfg["fire-hidden"],
            action_dim=cfg["fire-action"],
            actor_lr=cfg["actor-lr"],
            critic_lr=cfg["critic-lr"],
            alpha_lr=cfg['alpha-lr'],
            target_entropy=cfg['target-entropy'],
            tau=cfg['tau'],
            gamma=cfg['gamma'],
            device=cfg['device'],
        )
        self.water_model = SACContinuous(
            env,
            state_dim=cfg["state_dim"],
            hidden_dim=cfg["water-hidden"],
            action_dim=cfg["water-action"],
            actor_lr=cfg["actor-lr"],
            critic_lr=cfg["critic-lr"],
            alpha_lr=cfg['alpha-lr'],
            target_entropy=cfg['target-entropy'],
            tau=cfg['tau'],
            gamma=cfg['gamma'],
            device=cfg['device'],
        )
        self.new_energy = SACContinuous(
            env,
            state_dim=cfg["state_dim"],
            hidden_dim=cfg["energy-hidden"],
            action_dim=cfg["energy-action"],
            actor_lr=cfg["actor-lr"],
            critic_lr=cfg["critic-lr"],
            alpha_lr=cfg['alpha-lr'],
            target_entropy=cfg['target-entropy'],
            tau=cfg['tau'],
            gamma=cfg['gamma'],
            device=cfg['device'],
        )
        self.exchange = SACContinuous(
            env,
            state_dim=cfg["state_dim"],
            hidden_dim=cfg["exchage-hidden"],
            action_dim=cfg["exchage-action"],
            actor_lr=cfg["actor-lr"],
            critic_lr=cfg["critic-lr"],
            alpha_lr=cfg['alpha-lr'],
            target_entropy=cfg['target-entropy'],
            tau=cfg['tau'],
            gamma=cfg['gamma'],
            device=cfg['device'],
        )
    def take_action(self, state)->dict:
        # state is [curr_t, payload, battery_t, ...]
        
        # firstly take fire power action
        fire_action = self.fire_model.take_action(state)
        fire_power = fire_action.sum()
        # state[1] -= fire_power
        # if state[1] < 0: state[1] = 0
        water_action = self.water_model.take_action(state)
        water_power = water_action.sum()
        new_energy_action = self.new_energy.take_action(state)
        new_energy_power = new_energy_action.sum()
        exchange_action = self.exchange.take_action(state)
        exchange_power = exchange_action.sum()
        return {
            "fire": fire_action,
            "water": water_action,
            "energy": new_energy_action,
            "exchange": exchange_action
        }
    def update(self, ):
        pass
